import os
from enum import Enum
import time
import threading

import numpy as np
import mujoco
import mujoco.viewer
import glfw

import rospy
import rospkg

from std_msgs.msg import Float64MultiArray
from nav_msgs.msg import Odometry
from sensor_msgs.msg import Imu, JointState
from std_srvs.srv import SetBool, SetBoolRequest
from ocs2_msgs.msg import mpc_state, mpc_input, mpc_observation, mpc_target_trajectories, mode_schedule


class QuadrupedEnv:
    def __init__(self):
        super().__init__()
        self._robot_name = "legged_robot"
        rospy.init_node(self._robot_name + "_env", anonymous=True)
        package_path = rospkg.RosPack().get_path("mpc_quadruped_mujoco")
        xml_path = os.path.join(package_path, "assets", "anybotics_anymal_c", "scene.xml")
        self._mrt_desired_frequency = 1000

        self._state_dim = 24
        self._input_dim = 24
        self._actuated_dof_num = 12

        self._initial_position = [-0.25, 0.60, -0.85, -0.25, -0.60, 0.85, 0.25, 0.60, -0.85, 0.25, -0.60, 0.85]

        class ModeNumber(Enum):
            FLY = 0
            RH = 1
            LH = 2
            LH_RH = 3
            RF = 4
            RF_RH = 5
            RF_LH = 6
            RF_LH_RH = 7
            LF = 8
            LF_RH = 9
            LF_LH = 10
            LF_LH_RH = 11
            LF_RF = 12
            LF_RF_RH = 13
            LF_RF_LH = 14
            STANCE = 15

        stance = {
            "mode_sequence": [ModeNumber.STANCE.value],
            "switching_times": [0.0, 0.5]
        }
        trot = {
            "mode_sequence": [ModeNumber.LF_RH.value, ModeNumber.RF_LH.value],
            "switching_times": [0.0, 0.35, 0.70]
        }
        standing_trot = {
            "mode_sequence": [ModeNumber.LF_RH.value, ModeNumber.STANCE.value, ModeNumber.RF_LH.value,
                              ModeNumber.STANCE.value],
            "switching_times": [0.0, 0.30, 0.35, 0.65, 0.70]
        }
        flying_trot = {
            "mode_sequence": [ModeNumber.LF_RH.value, ModeNumber.FLY.value, ModeNumber.RF_LH.value,
                              ModeNumber.FLY.value],
            "switching_times": [0.0, 0.27, 0.30, 0.57, 0.60]
        }
        pace = {
            "mode_sequence": [ModeNumber.LF_LH.value, ModeNumber.FLY.value, ModeNumber.RF_RH.value,
                              ModeNumber.FLY.value],

            "switching_times": [0.0, 0.28, 0.30, 0.58, 0.60]
        }
        standing_pace = {
            "mode_sequence": [ModeNumber.LF_LH.value, ModeNumber.STANCE.value, ModeNumber.RF_RH.value,
                              ModeNumber.STANCE.value],
            "switching_times": [0.0, 0.30, 0.35, 0.65, 0.70]
        }
        dynamic_walk = {
            "mode_sequence": [ModeNumber.LF_RF_RH.value, ModeNumber.RF_RH.value, ModeNumber.RF_LH_RH.value,
                              ModeNumber.LF_RF_LH.value, ModeNumber.LF_LH.value, ModeNumber.LF_LH_RH.value],
            "switching_times": [0.0, 0.20, 0.30, 0.50, 0.70, 0.80, 1.0]
        }
        static_walk = {
            "mode_sequence": [ModeNumber.LF_RF_RH.value, ModeNumber.RF_LH_RH.value, ModeNumber.LF_RF_LH.value,
                              ModeNumber.LF_LH_RH.value],
            "switching_times": [0.0, 0.30, 0.60, 0.90, 1.20]
        }
        amble = {
            "mode_sequence": [ModeNumber.RF_LH.value, ModeNumber.LF_LH.value, ModeNumber.LF_RH.value,
                              ModeNumber.RF_RH.value],
            "switching_times": [0.0, 0.15, 0.40, 0.55, 0.80]
        }
        lindyhop = {
            "mode_sequence": [ModeNumber.LF_RH.value, ModeNumber.STANCE.value, ModeNumber.RF_LH.value,
                              ModeNumber.STANCE.value, ModeNumber.LF_LH.value, ModeNumber.RF_RH.value,
                              ModeNumber.LF_LH.value, ModeNumber.STANCE.value, ModeNumber.RF_RH.value,
                              ModeNumber.LF_LH.value, ModeNumber.RF_RH.value, ModeNumber.STANCE.value],
            "switching_times": [0.0, 0.35, 0.45, 0.80, 0.90, 1.125, 1.35, 1.70, 1.80, 2.025, 2.25, 2.60, 2.70]
        }
        skipping = {
            "mode_sequence": [ModeNumber.LF_RH.value, ModeNumber.FLY.value, ModeNumber.LF_RH.value,
                              ModeNumber.FLY.value, ModeNumber.RF_LH.value, ModeNumber.FLY.value,
                              ModeNumber.RF_LH.value, ModeNumber.FLY.value],
            "switching_times": [0.0, 0.27, 0.30, 0.57, 0.60, 0.87, 0.90, 1.17, 1.20]
        }
        pawup = {
            "mode_sequence": [ModeNumber.RF_LH_RH.value],
            "switching_times": [0.0, 2.0]
        }

        self._mode_schedule_dict = {
            "stance": stance,
            "trot": trot,
            "standing_trot": standing_trot,
            "flying_trot": flying_trot,
            "pace": pace,
            "standing_pace": standing_pace,
            "dynamic_walk": dynamic_walk,
            "static_walk": static_walk,
            "amble": amble,
            "lindyhop": lindyhop,
            "skipping": skipping,
            "pawup": pawup
        }

        self._odometry_msg = Odometry()
        self._imu_msg = Imu()
        self._joint_state_msg = JointState()

        self._mpc_ctrl_msg = Float64MultiArray()
        self._mpc_ctrl_msg.data = [0.0 for _ in range(self._actuated_dof_num * 3)]

        self._mpc_observation_msg = mpc_observation()

        self._kp = np.ones(self._actuated_dof_num) * 100
        self._kd = np.ones(self._actuated_dof_num) * 10

        self._received = False
        self._lock = threading.Lock()

        self._mj_model = mujoco.MjModel.from_xml_path(xml_path)
        self._mj_data = mujoco.MjData(self._mj_model)
        qpos = self._mj_data.qpos.copy()
        qpos[-self._actuated_dof_num:] = self._initial_position
        mujoco.mj_setState(self._mj_model, self._mj_data, qpos, mujoco.mjtState.mjSTATE_QPOS)
        mujoco.mj_forward(self._mj_model, self._mj_data)
        self._mj_viewer = mujoco.viewer.launch_passive(self._mj_model, self._mj_data, key_callback=self._key_callback)

        self._mpc_ctrl_subscriber = rospy.Subscriber(self._robot_name + "_mpc_ctrl", Float64MultiArray,
                                                     self._mpc_ctrl_callback)
        self._odometry_publisher = rospy.Publisher(self._robot_name + "_odometry", Odometry, queue_size=1)
        self._imu_publisher = rospy.Publisher(self._robot_name + "_imu", Imu, queue_size=1)
        self._joint_state_publisher = rospy.Publisher(self._robot_name + "_joint_state", JointState, queue_size=1)

        self._mpc_observation_publisher = rospy.Subscriber(self._robot_name + "_mpc_observation", mpc_observation,
                                                           self._mpc_observation_callback)
        self._target_trajectories_publisher = rospy.Publisher(self._robot_name + "_mpc_target", mpc_target_trajectories,
                                                              queue_size=1)
        self._mode_schedule_publisher = rospy.Publisher(self._robot_name + "_mpc_mode_schedule", mode_schedule,
                                                        queue_size=1)

        self._update_sensor()

        rospy.wait_for_service(self._robot_name + "_mpc_reset_mpc")
        self._mpc_reset_client = rospy.ServiceProxy(self._robot_name + "_mpc_reset_mpc", SetBool)

        self._mj_step_timer = rospy.Timer(rospy.Duration(1.0 / self._mrt_desired_frequency), self._mj_step_callback)
        self._mj_viewer_timer = rospy.Timer(rospy.Duration(0.02), self._mj_viewer_callback)

        request = SetBoolRequest()
        request.data = True
        response = self._mpc_reset_client.call(request)

        rospy.loginfo("---------- Target ----------")
        rospy.loginfo("+X:              Keypad 1")
        rospy.loginfo("-X:              Keypad 7")
        rospy.loginfo("+Y:              Keypad 6")
        rospy.loginfo("-Y:              Keypad 4")
        rospy.loginfo("+Z:              Keypad 8")
        rospy.loginfo("-Z:              Keypad 2")
        rospy.loginfo("+Yaw:            Keypad 9")
        rospy.loginfo("-Yaw:            Keypad 3")
        rospy.loginfo("Stop:            Keypad 5")

        rospy.loginfo("----------- Mode -----------")
        rospy.loginfo("Stance:          Shift + Q")
        rospy.loginfo("Trot:            Shift + W")
        rospy.loginfo("Standing Trot:   Shift + E")
        rospy.loginfo("Flying Trot:     Shift + R")
        rospy.loginfo("Pace:            Shift + A")
        rospy.loginfo("Standing Pace:   Shift + S")
        rospy.loginfo("Dynamic Walk:    Shift + D")
        rospy.loginfo("Static Walk:     Shift + F")
        rospy.loginfo("Amble:           Shift + Z")
        rospy.loginfo("Lindy Hop:       Shift + X")
        rospy.loginfo("Skipping:        Shift + C")
        rospy.loginfo("Paw Up:          Shift + V")

    def run(self):
        rospy.spin()

    def _update_ctrl(self):
        pos_desired = np.array(self._mpc_ctrl_msg.data[:self._actuated_dof_num])
        vel_desired = np.array(self._mpc_ctrl_msg.data[self._actuated_dof_num: self._actuated_dof_num * 2])
        tau_desired = np.array(self._mpc_ctrl_msg.data[self._actuated_dof_num * 2: self._actuated_dof_num * 3])

        pos_actual = self._mj_data.qpos[-self._actuated_dof_num:]
        vel_actual = self._mj_data.qvel[-self._actuated_dof_num:]

        self._mj_data.ctrl[:] = self._kp * (pos_desired - pos_actual) + self._kd * (
                vel_desired - vel_actual) + tau_desired

    def _key_callback(self, keycode):
        if glfw.KEY_KP_0 <= keycode <= glfw.KEY_KP_9:
            mpc_target_trajectories_msg = mpc_target_trajectories()
            mpc_target_trajectories_msg.timeTrajectory.extend(
                [self._mpc_observation_msg.time, self._mpc_observation_msg.time + 1.0])

            mpc_input1 = mpc_input()
            mpc_input2 = mpc_input()
            mpc_input1.value.extend([0.0 for _ in range(self._input_dim)])
            mpc_input2.value.extend([0.0 for _ in range(self._input_dim)])
            mpc_target_trajectories_msg.inputTrajectory.extend([mpc_input1, mpc_input2])

            mpc_state1 = mpc_state()
            mpc_state2 = mpc_state()

            mpc_state1.value = [0.0 for _ in range(self._state_dim)]
            mpc_state2.value = [0.0 for _ in range(self._state_dim)]

            mpc_state1.value[6:12] = self._mpc_observation_msg.state.value[6:12]
            mpc_state1.value[12: 12 + self._actuated_dof_num] = self._initial_position
            mpc_state2.value[6:12] = self._mpc_observation_msg.state.value[6:12]
            mpc_state2.value[12: 12 + self._actuated_dof_num] = self._initial_position

            if keycode == glfw.KEY_KP_1:
                mpc_state2.value[6] += 0.5
            elif keycode == glfw.KEY_KP_7:
                mpc_state2.value[6] -= 0.5
            elif keycode == glfw.KEY_KP_6:
                mpc_state2.value[7] += 0.5
            elif keycode == glfw.KEY_KP_4:
                mpc_state2.value[7] -= 0.5
            elif keycode == glfw.KEY_KP_8:
                mpc_state2.value[8] += 0.5
            elif keycode == glfw.KEY_KP_2:
                mpc_state2.value[8] -= 0.5
            elif keycode == glfw.KEY_KP_9:
                mpc_state2.value[9] += 0.3
            elif keycode == glfw.KEY_KP_3:
                mpc_state2.value[9] -= 0.3
            elif keycode == glfw.KEY_KP_5:
                pass
            else:
                return
            mpc_target_trajectories_msg.stateTrajectory.extend([mpc_state1, mpc_state2])
            self._target_trajectories_publisher.publish(mpc_target_trajectories_msg)

        if glfw.KEY_A <= keycode <= glfw.KEY_Z:
            mode_schedule_msg = mode_schedule()
            mode = ""
            if keycode == glfw.KEY_Q:
                mode = "stance"
            elif keycode == glfw.KEY_W:
                mode = "trot"
            elif keycode == glfw.KEY_E:
                mode = "standing_trot"
            elif keycode == glfw.KEY_R:
                mode = "flying_trot"
            elif keycode == glfw.KEY_A:
                mode = "pace"
            elif keycode == glfw.KEY_S:
                mode = "standing_pace"
            elif keycode == glfw.KEY_D:
                mode = "dynamic_walk"
            elif keycode == glfw.KEY_F:
                mode = "static_walk"
            elif keycode == glfw.KEY_Z:
                mode = "amble"
            elif keycode == glfw.KEY_X:
                mode = "lindyhop"
            elif keycode == glfw.KEY_C:
                mode = "skipping"
            elif keycode == glfw.KEY_V:
                mode = "pawup"
            else:
                return
            mode_schedule_msg.eventTimes.extend(self._mode_schedule_dict[mode]["switching_times"])
            mode_schedule_msg.modeSequence.extend(self._mode_schedule_dict[mode]["mode_sequence"])
            self._mode_schedule_publisher.publish(mode_schedule_msg)
            return

    def _update_sensor(self):
        mujoco.mj_forward(self._mj_model, self._mj_data)
        pos = self._mj_data.sensor('BodyPos').data.copy()
        quat = self._mj_data.sensor('BodyQuat').data.copy()
        vel = self._mj_data.sensor('BodyVel').data.copy()
        angVel = self._mj_data.sensor('BodyGyro').data.copy()
        acc = self._mj_data.sensor('BodyAcc').data.copy()

        self._odometry_msg.header.stamp = rospy.Time.now()
        self._odometry_msg.pose.pose.position.x = pos[0]
        self._odometry_msg.pose.pose.position.y = pos[1]
        self._odometry_msg.pose.pose.position.z = pos[2]
        self._odometry_msg.pose.pose.orientation.w = quat[0]
        self._odometry_msg.pose.pose.orientation.x = quat[1]
        self._odometry_msg.pose.pose.orientation.y = quat[2]
        self._odometry_msg.pose.pose.orientation.z = quat[3]
        self._odometry_msg.twist.twist.linear.x = vel[0]
        self._odometry_msg.twist.twist.linear.y = vel[1]
        self._odometry_msg.twist.twist.linear.z = vel[2]
        self._odometry_msg.twist.twist.angular.x = angVel[0]
        self._odometry_msg.twist.twist.angular.y = angVel[1]
        self._odometry_msg.twist.twist.angular.z = angVel[2]

        self._imu_msg.header.stamp = rospy.Time.now()
        self._imu_msg.orientation.w = quat[0]
        self._imu_msg.orientation.x = quat[1]
        self._imu_msg.orientation.y = quat[2]
        self._imu_msg.orientation.z = quat[3]
        self._imu_msg.angular_velocity.x = angVel[0]
        self._imu_msg.angular_velocity.y = angVel[1]
        self._imu_msg.angular_velocity.z = angVel[2]
        self._imu_msg.linear_acceleration.x = acc[0]
        self._imu_msg.linear_acceleration.y = acc[1]
        self._imu_msg.linear_acceleration.z = acc[2]
        self._imu_msg.orientation_covariance = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        self._imu_msg.angular_velocity_covariance = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        self._imu_msg.linear_acceleration_covariance = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]

        self._joint_state_msg.position = self._mj_data.qpos[-self._actuated_dof_num:].tolist()
        self._joint_state_msg.velocity = self._mj_data.qvel[-self._actuated_dof_num:].tolist()

    def _publish_sensor(self):
        self._odometry_publisher.publish(self._odometry_msg)
        self._imu_publisher.publish(self._imu_msg)
        self._joint_state_publisher.publish(self._joint_state_msg)

    def _mpc_ctrl_callback(self, msg: Float64MultiArray):
        self._mpc_ctrl_msg.data = msg.data
        self._received = True

    def _mpc_observation_callback(self, msg: mpc_observation):
        self._mpc_observation_msg = msg

    def _mj_step_callback(self, event):
        if not self._received:
            return

        self._lock.acquire()
        self._update_ctrl()
        mujoco.mj_step(self._mj_model, self._mj_data)
        self._update_sensor()
        self._lock.release()
        self._publish_sensor()

    def _mj_viewer_callback(self, event):
        self._lock.acquire()
        self._mj_viewer.sync()
        self._lock.release()


if __name__ == '__main__':
    node = QuadrupedEnv()
    node.run()
